import torch
from PIL import Image, ImageFilter
import os
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, img_path, txt_path, transform):
        super(MyDataset, self).__init__()
        self.img_root = img_path
        self.transform = transform
        self.fnames = []
        self.txt = []
        self.labels = []

        for txt_name in os.listdir(txt_path):
            img_name = txt_name[:-4] + '.jpg'
            self.fnames.append(img_name)

            result = open(os.path.join(txt_path, txt_name))
            result = result.read().split()
            self.txt.append(torch.Tensor([float(result[0])/200, float(result[1])/100, float(result[2])/600, float(result[3])/0.25]))
            self.labels.append(torch.Tensor([[float(s)/180 for s in result[4:]]]).mean(-1))
        self.samples = len(self.fnames)

    def __getitem__(self, item):
        fname = self.fnames[item]
        img = Image.open(os.path.join(self.img_root, fname))
        txt = self.txt[item].clone()
        label = self.labels[item].clone()
        img = self.transform(img)

        return img, txt, label

    def __len__(self):
        return self.samples


if __name__ == '__main__':
    import torchvision.transforms as transforms
    from torch.utils.data import DataLoader
    from config import *
    from model import ContactAngle

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])

    model = ContactAngle(out_puts=2)

    train_data = MyDataset(img_path=img_root,
                           txt_path=txt_root, transform=transform)

    train_data_loader = DataLoader(dataset=train_data, batch_size=10, shuffle=True)

    for i, (img, txt, label) in enumerate(train_data_loader):
        out = model(img, txt)
        print(out)
